from langchain.agents import load_tools, initialize_agent, AgentType
from langchain.callbacks import get_openai_callback
from langchain.llms.openai import OpenAI

llm = OpenAI(model_name="gpt-3.5-turbo-instruct", n=2, best_of=2)
with get_openai_callback() as cb:
    result = llm.invoke("Tell me a joke")
    print(cb)
    print(cb.total_tokens)


def step_usage():
    llm = OpenAI(temperature=0)
    tools = load_tools(["serpapi", "llm-math"], llm)
    agent = initialize_agent(
        tools, llm, AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
    )
    with get_openai_callback() as cb:
        response = agent.run("Who is Olivia Wilde's boyfriend? What is his current age raised to the 0.23 power?")
        print(f"Total Tokens: {cb.total_tokens}")
        print(f"Prompt Tokens: {cb.prompt_tokens}")
        print(f"Completion Tokens: {cb.completion_tokens}")
        print(f"Total Cost (USD): ${cb.total_cost}")


if __name__ == '__main__':
    step_usage()
